import os
from glob import glob
from tqdm import tqdm
from shutil import copy


def get_img_path(img_dir, suffix= ("jpg", "png")):
    img_paths = []
    for sfx in suffix:
        for path in glob(os.path.join(img_dir, f"*.{sfx}")):
            img_paths.append(path)

    return img_paths


def main():
    img_dir = r"D:\lliujian\DeepLearning\datasets\FPC_20250301_ori\images\train"
    txt_dir = r"D:\lliujian\DeepLearning\datasets\FPC_20250301\FPC_20250301\labels\train"
    save_txt_dir = r"D:\lliujian\DeepLearning\datasets\FPC_20250301_ori\labels\train"
    for img_path in tqdm(get_img_path(img_dir)):
        img_name = os.path.basename(img_path)
        suffix = os.path.splitext(img_path)[1]
        txt_path = os.path.join(txt_dir, img_name.replace(suffix, ".txt"))
        save_txt_path = os.path.join(save_txt_dir, img_name.replace(suffix, ".txt"))
        if os.path.exists(txt_path):
            copy(txt_path, os.path.join(img_dir, save_txt_path))
        else:
            print(f"{txt_path} not exists!")


if __name__ == '__main__':
    main()